package mws;

import java.util.concurrent.atomic.AtomicLong;

public class SlidingWindowRateLimiter {
    private final int windowSize; // 时间窗口大小，单位毫秒
    private final int segmentSize; // 每个格子的大小，单位毫秒
    private final int segmentCount; // 格子数量
    private final long limit; // 限流阈值

    private final AtomicLong[] counters; // 计数器数组
    private long lastUpdateTime; // 最后更新时间
    private long windowStart; // 当前窗口的开始时间

    public SlidingWindowRateLimiter(int windowSize, int segmentCount, long limit) {
        this.windowSize = windowSize;
        this.segmentCount = segmentCount;
        this.segmentSize = windowSize / segmentCount;
        this.limit = limit;
        this.counters = new AtomicLong[segmentCount];
        for (int i = 0; i < segmentCount; i++) {
            counters[i] = new AtomicLong(0);
        }
        this.lastUpdateTime = System.currentTimeMillis();
        this.windowStart = System.currentTimeMillis();
    }

    public boolean tryAcquire() {
        return tryAcquire(1);
    }

    public boolean tryAcquire(int permits) {
        long currentTime = System.currentTimeMillis();
        synchronized (this) {
            // 计算当前时间与最后更新时间的时间差，淘汰过期的格子
            long timePassed = currentTime - lastUpdateTime;
            if (timePassed >= windowSize) {
                // 如果已经过了一个窗口时间，重置所有计数器
                resetCounters();
                windowStart = currentTime;
            } else {
                // 淘汰过期的格子，将过期的格子计数器清零
                long segmentsToPass = timePassed / segmentSize;
                if (segmentsToPass > 0) {
                    // 计算需要清零的格子范围
                    int startIndex = (int) ((lastUpdateTime - windowStart) / segmentSize);
                    for (long i = 0; i < segmentsToPass; i++) {
                        int index = (int) ((startIndex + i) % segmentCount);
                        counters[index].set(0);
                    }
                    // 更新窗口开始时间，实际上窗口开始时间应该是当前时间减去窗口大小，但这里我们通过格子清零来模拟滑动
                    windowStart += segmentsToPass * segmentSize;
                }
            }
            lastUpdateTime = currentTime;

            // 计算当前窗口内的请求总数
            long total = 0;
            for (AtomicLong counter : counters) {
                total += counter.get();
            }

            // 如果超过限流阈值，则拒绝
            if (total + permits > limit) {
                return false;
            }

            // 找到当前时间对应的格子，增加计数
            int currentIndex = (int) (((currentTime - windowStart) / segmentSize) % segmentCount);
            counters[currentIndex].addAndGet(permits);
            return true;
        }
    }

    private void resetCounters() {
        for (AtomicLong counter : counters) {
            counter.set(0);
        }
    }
}